package com.zhaorujie.rpc.core.register.impl;

import com.zhaorujie.rpc.core.exception.RpcException;
import com.zhaorujie.rpc.core.register.Register;
import com.zhaorujie.rpc.core.utils.IpUtil;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;

import java.util.*;

/**
 * @author 赵汝杰
 * @version v1.0
 * @description ZookeeperRegister
 * @date 2024/3/5
 **/
public class ZookeeperRegister implements Register {
    public static final int BASE_SLEEP_TIME_MS = 1000;
    public static final String ZK_BASE_PATH = "/netty_rpc";
    public static final int MAX_RETRIES = 3;

    private ServiceDiscovery<String> serviceDiscovery;

    @Override
    public void init(String registerAddr, String env) {
        CuratorFramework client = CuratorFrameworkFactory.newClient(registerAddr, new ExponentialBackoffRetry(BASE_SLEEP_TIME_MS, MAX_RETRIES));
        client.start();

        JsonInstanceSerializer<String> instanceSerializer = new JsonInstanceSerializer<>(String.class);
        this.serviceDiscovery = ServiceDiscoveryBuilder.builder(String.class)
                .client(client)
                .basePath(Objects.isNull(env) || Objects.equals(env, "") ? ZK_BASE_PATH : ZK_BASE_PATH + "/" + env)
                .serializer(instanceSerializer)
                .build();
        try {
            this.serviceDiscovery.start();
        } catch (Exception e) {
            throw new RpcException(e);
        }
    }

    @Override
    public void register(String serviceKey, String serviceAddr) {
        try {
            String[] address = IpUtil.parseIpPort(serviceAddr);
            ServiceInstance<String> serviceInstance = ServiceInstance.<String>builder()
                    .name(serviceKey)
                    .address(address[0])
                    .port(Integer.parseInt(address[1]))
                    .payload(serviceAddr)
                    .build();
            serviceDiscovery.registerService(serviceInstance);
        } catch (Exception e) {
            throw new RpcException(e);
        }
    }

    @Override
    public void register(Collection<String> serviceKeys, String serviceAddr) {
        for (String serviceKey : serviceKeys) {
            register(serviceKey, serviceAddr);
        }
    }

    @Override
    public List<String> discovery(String serviceName) {
        try {
            List<String> list = new ArrayList<>();
            Collection<ServiceInstance<String>> serviceInstances = serviceDiscovery.queryForInstances(serviceName);
            for (ServiceInstance<String> serviceInstance : serviceInstances) {
                list.add(serviceInstance.getPayload());
            }
            return list;
        } catch (Exception e) {
            throw new RpcException(e);
        }

    }

    @Override
    public Map<String, List<String>> discovery(Collection<String> serviceKeys) {
        Map<String, List<String>> registryDataMap = new HashMap<>();
        for (String key : serviceKeys) {
            List<String> discovery = discovery(key);
            registryDataMap.put(key, discovery);
        }
        return registryDataMap;
    }

    @Override
    public void removeRegister(String serviceKey, String serviceAddr) {
        try {
            String[] address = IpUtil.parseIpPort(serviceAddr);
            ServiceInstance<String> serviceInstance = ServiceInstance.<String>builder()
                    .name(serviceKey)
                    .address(address[0])
                    .port(Integer.parseInt(address[1]))
                    .payload(serviceAddr)
                    .build();
            serviceDiscovery.unregisterService(serviceInstance);
        } catch (Exception e) {
            throw new RpcException(e);
        }
    }

    @Override
    public void removeRegister(Collection<String> serviceKeys, String serviceAddr) {
        for (String serviceKey : serviceKeys) {
            removeRegister(serviceKey, serviceAddr);
        }
    }

    @Override
    public void destroy() {

    }
}
